#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
检索器初始化测试

测试各个检索器的初始化功能
"""

import sys
import os
from loguru import logger

# 添加项目根目录到路径
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

# 在导入任何库之前强制设置环境变量
os.environ["HF_ENDPOINT"] = "https://hf-mirror.com"
os.environ["HUGGINGFACE_CO_URL_HOME"] = "https://hf-mirror.com"
os.environ["HUGGINGFACE_HUB_CACHE"] = os.path.expanduser("~/.cache/huggingface")
os.environ["SENTENCE_TRANSFORMERS_HOME"] = os.path.expanduser("~/.cache/sentence_transformers")

# 强制导入配置，确保镜像环境变量已设置
import questionretrieval.config

logger.info("🌐 使用 HuggingFace 国内镜像配置")
logger.info(f"📁 缓存目录: {questionretrieval.config.HF_HUB_CACHE}")
logger.info(f"🔗 镜像地址: {questionretrieval.config.HF_MIRROR_ENDPOINT}")

import pytest

def test_sentence_transformer_init():
    """
    测试 SentenceTransformer 检索器初始化
    """
    logger.info("\n=== 测试 SentenceTransformer 初始化 ===")
    
    try:
        from questionretrieval import SentenceTransformerRetriever
        retriever = SentenceTransformerRetriever()
        logger.info("✓ SentenceTransformer 初始化成功")
        return True
    except ImportError:
        logger.warning("✗ SentenceTransformer 不可用 (缺少依赖)")
        return False
    except Exception as e:
        logger.error(f"✗ SentenceTransformer 初始化失败: {e}")
        return False

def test_bert_chinese_init():
    """
    测试 BERT 中文检索器初始化
    """
    logger.info("\n=== 测试 BERT 中文检索器初始化 ===")
    
    try:
        from questionretrieval import BertChineseRetriever
        retriever = BertChineseRetriever()
        logger.info("✓ BERT 中文检索器初始化成功")
        return True
    except ImportError:
        logger.warning("✗ BertChineseRetriever 不可用 (缺少依赖)")
        return False
    except Exception as e:
        logger.error(f"✗ BERT 中文检索器初始化失败: {e}")
        return False

def test_bert_similarity_init():
    """
    测试 BERT 相似度检索器初始化
    """
    logger.info("\n=== 测试 BERT 相似度检索器初始化 ===")
    
    try:
        from questionretrieval import BertSimilarityRetriever
        retriever = BertSimilarityRetriever()
        logger.info("✓ BERT 相似度检索器初始化成功")
        return True
    except ImportError:
        logger.warning("✗ BertSimilarityRetriever 不可用 (缺少依赖)")
        return False
    except Exception as e:
        logger.error(f"✗ BERT 相似度检索器初始化失败: {e}")
        return False

def test_dbscan_clustering_init():
    """
    测试 DBSCAN 聚类检索器初始化
    """
    logger.info("\n=== 测试 DBSCAN 聚类检索器初始化 ===")
    
    try:
        from questionretrieval import DBSCANClusteringRetriever
        retriever = DBSCANClusteringRetriever()
        logger.info("✓ DBSCAN 聚类检索器初始化成功")
        return True
    except ImportError:
        logger.warning("✗ DBSCANClusteringRetriever 不可用 (缺少依赖)")
        return False
    except Exception as e:
        logger.error(f"✗ DBSCAN 聚类检索器初始化失败: {e}")
        return False

def test_hierarchical_clustering_init():
    """
    测试层级聚类检索器初始化
    """
    logger.info("\n=== 测试层级聚类检索器初始化 ===")
    
    try:
        from questionretrieval import HierarchicalClusteringRetriever
        retriever = HierarchicalClusteringRetriever()
        logger.info("✓ 层级聚类检索器初始化成功")
        return True
    except ImportError:
        logger.warning("✗ HierarchicalClusteringRetriever 不可用 (缺少依赖)")
        return False
    except Exception as e:
        logger.error(f"✗ 层级聚类检索器初始化失败: {e}")
        return False

def test_all_retrievers_init():
    """
    测试所有检索器初始化
    """
    logger.info("\n" + "="*50)
    logger.info("检索器初始化测试")
    logger.info("="*50)
    
    results = {
        'SentenceTransformer': test_sentence_transformer_init(),
        'BertChinese': test_bert_chinese_init(),
        'BertSimilarity': test_bert_similarity_init(),
        'DBSCANClustering': test_dbscan_clustering_init(),
        'HierarchicalClustering': test_hierarchical_clustering_init()
    }
    
    logger.info("\n" + "-"*40)
    logger.info("初始化测试结果")
    logger.info("-"*40)
    
    success_count = 0
    for name, success in results.items():
        status = "✓ 成功" if success else "✗ 失败"
        logger.info(f"{name:20} : {status}")
        if success:
            success_count += 1
    
    logger.info(f"\n总计: {success_count}/{len(results)} 个检索器可以初始化")
    
    if success_count > 0:
        logger.info("\n💡 提示: 可以初始化的检索器可以进行后续测试")
    else:
        logger.warning("\n⚠️  没有检索器可以初始化，请检查依赖安装")
        logger.info("解决方案: pip install transformers sentence-transformers torch scikit-learn")
    
    return results

if __name__ == "__main__":
    test_all_retrievers_init()